from __future__ import division
from xml.etree import cElementTree
import sys, re, random, math, cPickle
from util import *
from collections import defaultdict

class MNBTraining():
    def __init__(self, collectionfile, trainfile):
        self.cfile = collectionfile
        self.docnum = 0
        self.tags = ['0','1','2','3','4','5','6']
        self.titlewordcount = []
        self.textwordcount = []
        self.prior = dict()
        self.titlevocabulary = set()
        self.textvocabulary = set()
        self.titlecondprob = dict()
        self.textcondprob = dict()
        self.selected = 9353
        t1 = now()
        self.extractvocabulary()
        t2 = now()
        print t2-t1
        print 'finish extracting vocabulary.'

        t1 = now()        
        self.train(trainfile)
        t2 = now()
        print t2 - t1
        print 'finish tranning the data.'

        self.output()

    def train(self, trainfile):
        tag = dict()
        count = 0

        for line in open(trainfile, 'r'):
            temp = re.sub('\n','', line).split(' ')
            count += 1
            if tag.has_key(temp[1]):
                tag[temp[1]].append(int(temp[0]))
            else:
                tag[temp[1]] = [int(temp[0])]

        for key in tag:
            titlecondprob_tag = defaultdict(int)
            textcondprob_tag = defaultdict(int)
            N1 = 0
            N2 = 0

            for doc in tag[key]:
                N1 += sum(self.titlewordcount[doc].values())
                N2 += sum(self.textwordcount[doc].values())
                for word in self.titlevocabulary:
                    titlecondprob_tag[word] += self.titlewordcount[doc].get(word, 0)
                for word in self.textvocabulary:
                    textcondprob_tag[word] += self.textwordcount[doc].get(word, 0)
            for i in titlecondprob_tag:
                titlecondprob_tag[i] = (titlecondprob_tag[i] + 1) / (N1 + len(self.titlevocabulary))
            for i in textcondprob_tag:
                textcondprob_tag[i] = (textcondprob_tag[i] + 1) / (N2 + len(self.textvocabulary))

            self.titlecondprob[key] = titlecondprob_tag
            self.textcondprob[key] = textcondprob_tag
            self.prior[key] = len(tag[key]) / count



    def output(self):
        f = open('labMNB.txt', 'w')
        for i in range(0, self.docnum):
            test = 0
            score = dict()

            for t in self.tags:
                score[t] = math.log10(self.prior[t])
                for word in self.titlewordcount[i]:
                    score[t] += 0.7*math.log10(self.titlecondprob[t][word])
                for word in self.textwordcount[i]:
                    score[t] += 0.3*math.log10(self.textcondprob[t][word])

            f.write(argmax(score)+'\n')
        f.close()
        '''save prior and conditional pro into a file'''

    def extractvocabulary(self):
        """
        try1: stem
        try2: no stem
        try3: log10
        try4: weight title
        """
        text = ''
        docid = 0
        weight = 1
        for event, elem in cElementTree.iterparse(self.cfile):
            count = dict()
            t = re.sub('{.*}','',elem.tag.lower())
            if t == 'page':
                docid += 1
            elif t == 'text':
                tokens = re.split('[^A-Za-z]+', elem.text.lower())
                for i in tokens:
                    token = stem(i)
                    '''token = i'''
                    if not token: continue
                    elif token in STOPWORDS: continue
                    else:
                        if count.has_key(token):
                            count[token] += 1
                        else:
                            self.textvocabulary.add(token)
                            count[token] = 1
                self.textwordcount.append(count)
            elif t == 'title': 
                tokens = re.split('[^A-Za-z]+', elem.text.lower())

                for i in tokens:
                    '''token = stem(i)'''
                    token = i
                    if not token: continue
                    elif token in STOPWORDS: continue
                    else:
                        if count.has_key(token):
                            count[token] += 1
                        else:
                            self.titlevocabulary.add(token)
                            count[token] = 1
                self.titlewordcount.append(count)
            elem.clear()
        self.docnum = docid


if __name__ == "__main__":
    arg = sys.argv[1:]
    if len(arg) != 2:
        print ''''''
    collectionfile = arg[0]
    trainfile = arg[1]
    try:
        import psyco
        psyco.full()
    except:
        pass
    bt = MNBTraining(collectionfile, trainfile)
    